import json

from flask import Blueprint, request, Response, stream_with_context

from llm.llm import get_llm_chain
from model.result_model import Result

api_bp = Blueprint('api', __name__)



@api_bp.route('/stream', methods=['POST'])
def ai_stream():
    question = request.json.get('question')
    session_id = ''
    resp = get_llm_chain("prompt").invoke({"input": question, "chat_history": [], "agent_scratchpad": []},
                                           config={"configurable": {"session_id": session_id}},
                                           return_only_outputs=True)
    con = resp.content

    result = Result(
        code=200,
        msg="操作成功",
        data=con
    )
    json_str = json.dumps(result.to_dict(), ensure_ascii=False)
    return Response(
        stream_with_context(event_stream(con)),  # 流式上下文确保生成器正常工作
        mimetype='text/event-stream',  # 设置 MIME 类型为 SSE
        headers={
            'Cache-Control': 'no-cache',  # 禁用缓存
            'Connection': 'keep-alive'  # 保持连接
        }
    )

@api_bp.route('/chat', methods=['POST'])
def ai_chat():
    question = request.json.get('question')
    session_id = ''
    resp = get_llm_chain("prompt").invoke({"input": question, "chat_history": [], "agent_scratchpad": []},
                                           config={"configurable": {"session_id": session_id}},
                                           return_only_outputs=True)
    con = resp.content

    result = Result(
        code=200,
        msg="操作成功",
        data=con
    )
    json_str = json.dumps(result.to_dict(), ensure_ascii=False)
    return Response(json_str, mimetype='application/json; charset=utf-8')

def event_stream(s):
    """
    生成器函数，逐步返回事件流数据。
    """
    # for c in s:
    #     data = {'content': f'{c}'}
    #     yield f'data: {json.dumps(data)}\n\n'
    chunk_size = 50
    for i in range(0, len(s), chunk_size):
        chunk = s[i:i + chunk_size]
        data = {'content': chunk}
        #time.sleep(1)
        yield f'data: {json.dumps(data, ensure_ascii=False)}\n\n'
    # data = {'content': f'成功'}
    # yield f'data: {json.dumps(data)}\n\n'

    yield "event: end\ndata: [DONE]\n\n"